﻿using CPAIS.Alglorithm;
using CPAIS.Model;
using CPAIS.Updater;
using System;
using System.Collections.Generic;

namespace CPAIS.Adapter
{
    public class CPMGibbsAdapter
    {
        public struct TopicItem
        {
            public struct Item : IComparable<Item>
            {
                public struct TimeItem
                {
                    public int Day;
                    public double Prob;
                    public int Count;
                }

                public string ActionType;
                public double ActionProb;
                public TimeItem[] TimeProbs;

                public int CompareTo(Item other)
                {
                    if (this.ActionProb > other.ActionProb)
                        return -1;
                    else if (this.ActionProb < other.ActionProb)
                        return 1;
                    else
                        return 0;
                }
            }

            public double Prob;
            public Item[] Actions;
        }

        private readonly List<Trace> traces;
        private int M, V, T, K;
        private double alpha, beta, gamma;
        private int iterations;
        private readonly string[] vocab;
        private readonly int[] timelib;
        private readonly int[][] docs;
        private readonly int[][] times;
        private Dictionary<string, double[]> docTopicAssoc;
        private TopicItem[] topicWordAssoc;
        private double perplexity;
        private double klDivergence;
        private double entropy;

        public CPMGibbsAdapter(LogModel logModel)
        {
            traces = logModel.Traces;
            vocab = logModel.ActionTypes;
            timelib = logModel.ActionDays;
            
            M = traces.Count;
            V = vocab.Length;
            T = timelib.Length;

            docs = new int[M][];
            times = new int[M][];
            for (int m = 0; m < M; m++)
            {
                var actions = traces[m].Actions;
                int N = actions.Count;
                docs[m] = new int[N];
                times[m] = new int[N];
                for (int n = 0; n < N; n++)
                {
                    docs[m][n] = actions[n].TypeIndex;
                    int day = actions[n].Day;
                    times[m][n] = Array.IndexOf(timelib, day);
                }
            }
        }

        public int VocabularySize { get { return V; } }
        public int DocumentSize { get { return M; } }
        public int TimeLibrarySize { get { return T; } }
        public int TopicSize { get { return K; } }
        public double Alpha { get { return alpha; } }
        public double Beta { get { return beta; } }
        public double Gamma { get { return gamma; } }
        public int Iterations { get { return iterations; } }

        public Dictionary<string, double[]> DocTopicAssoc
        {
            get { return docTopicAssoc; }
        }
        public TopicItem[] TopicWordAssoc
        {
            get { return topicWordAssoc; }
        }
        public double Perplexity { get { return perplexity; } }
        public double KLDivergence { get { return klDivergence; } }
        public double Entropy { get { return entropy; } }

        public void Execute(int K, double alpha, double beta, double gamma, int iterations, ProgressUpdater updater)
        {
            // Records parameters
            this.K = K;
            this.alpha = alpha;
            this.beta = beta;
            this.gamma = gamma;
            this.iterations = iterations;

            // Generates parameters
            CPMGibbs.Params param = new CPMGibbs.Params();
            param.K = K;
            param.alpha = alpha;
            param.beta = beta;
            param.gamma = gamma;
            param.iterations = iterations;
            param.docs = docs;
            param.times = times;
            param.V = V;
            param.T = T;

            // Gibbs sampling
            CPMGibbs gibbs = new CPMGibbs(param);

            try
            {
                if (updater != null)
                    updater.Begin();

                gibbs.GibbsSampling(updater);

                // Gets Document-Topic association
                double[,] theta = gibbs.Theta;
                docTopicAssoc = new Dictionary<string, double[]>();
                for (int m = 0; m < M; m++)
                {
                    double[] probs = new double[K];
                    for (int k = 0; k < K; k++)
                    {
                        probs[k] = theta[m, k];
                    }
                    docTopicAssoc.Add(traces[m].ID, probs);
                }

                // Gets Topic-Word association
                double[,] phi = gibbs.Phi;
                double[, ,] psi = gibbs.Psi;
                int[, ,] psiCount = gibbs.PsiCount;
                topicWordAssoc = new TopicItem[K];
                for (int k = 0; k < K; k++)
                {
                    topicWordAssoc[k].Actions = new TopicItem.Item[V];
                    for (int v = 0; v < V; v++)
                    {
                        topicWordAssoc[k].Actions[v].ActionType = vocab[v];
                        topicWordAssoc[k].Actions[v].ActionProb = phi[k, v];
                        var timeProbs = new TopicItem.Item.TimeItem[T];
                        for (int t = 0; t < T; t++)
                        {
                            timeProbs[t].Day = timelib[t];
                            timeProbs[t].Prob = psi[k, v, t];
                            timeProbs[t].Count = psiCount[k, v, t];
                        }
                        topicWordAssoc[k].Actions[v].TimeProbs = timeProbs;
                    }
                    topicWordAssoc[k].Prob = 0;
                    foreach (var key in docTopicAssoc.Keys)
                    {
                        topicWordAssoc[k].Prob += docTopicAssoc[key][k];
                    }
                    topicWordAssoc[k].Prob /= M;
                    Array.Sort(topicWordAssoc[k].Actions);
                }

                // Gets other result
                perplexity = gibbs.CalcPerplexity();
                klDivergence = gibbs.CalcKLDivergence();
                entropy = gibbs.CalcEntropy();

                if (updater != null)
                {
                    updater.Complete();
                }
            }
            catch (Exception ex)
            {
                if (updater != null)
                {
                    if (updater.State == UpdateState.Cancelled)
                        return;
                    else
                        updater.Fail(ex.Message);
                }
                else
                    throw ex;

                return;
            }
        }
    }
}
